from pydantic import BaseModel
from typing import List, Optional, Union



class TokensCountRequestModel(BaseModel):
    model: str
    input: List[str]


class CompletionsRequestModel(BaseModel):
    model: str
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    n: Optional[int] = 1
    stream: Optional[bool] = False
    max_tokens: Optional[int] = 1024
    repetion_penalty: Optional[float] = 1.0
    update_interval: Optional[int] = None



    
    class Message(BaseModel):
        role: str | None
        content: str
        data_for_context: List
    

    class FunctionCall(BaseModel):
        name: str
        parameters: dict


    class Function(BaseModel):
        name:str
        descrtiption: Optional[str] = None
        parameters: dict


        class ShotExample(BaseModel):
            request: str
            params: dict

        few_shot_example: Optional[List[ShotExample]] = None
        return_parameters: Optional[dict] = None
        
    messages: List[Message]
    function_call: Union[FunctionCall, str]


class CreateEmbeddingRequestModel(BaseModel):
    model: str = 'Embedding'
    input: List[str]